{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, TypeFamilies,
  UndecidableInstances #-}

module LLBasicsEx where

import T

import qualified Data.IntMap as IM; import Data.IntMap (IntMap)

import Control.Monad (liftM, ap)
import Control.Applicative (Applicative(pure, (<*>)))





type Rename = (Int, IntMap Int)

lookupRN :: Rename -> Int -> Occ
lookupRN rn@(locals, m) i
  | i < locals = Par i
  | otherwise = maybe err Env $ IM.lookup (i - locals) m
  where err = error $ "LLBasics.lookupRN: unresolved variable: " ++ show (i, rn)



newtype M a =
  M {runM :: (Rename, Int) -> (a, [Dec])}

instance Functor M where fmap = liftM
instance Applicative M where pure = return; (<*>) = ap
instance Monad M where
  return a = M $ \_ -> (a, [])
  m >>= k = M $ \(rn, sh) ->
     -- NB backwards state: a and w' are circular
     let (a, w)  = runM m     (rn, sh + length w')
         (b, w') = runM (k a) (rn, sh)
     in (b, w ++ w')

ask :: M Rename
ask = M $ \ ~(rn, _) -> (rn, [])

local :: (Rename -> Rename) -> M a -> M a
local f (M g) = M $ \ ~(x, y) -> g (f x, y)

numEmissions :: M Int
numEmissions = M $ \ ~(_, s) -> (s, [])

resetEmissions :: M a -> M a
resetEmissions (M f) = M $ \(rn, _) -> f (rn, 0)

emit :: Dec -> M ()
emit w = M $ \_ -> ((), [w])
